package org.iscas.tcse.preprocessing;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Transaction;

import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

public class MaxAbsScaler extends ProcessFunction<DenseVector, DenseVector>
        implements Preprocessing {
    private static final Logger LOG = LoggerFactory.getLogger(StdScaler.class);

    public final String prefix;

    private final int windowSize;
    private int currSize = 0;
    private final int[] featureIndices;
    private volatile boolean bootWindow = true;

    private volatile double[] maxAbsValues;

    private volatile boolean isRunning = true;

    private final String[] keys;

    private final BlockingQueue<double[]> channel;
    private final String host;
    private final int port;

    private final String setMaxAbsScript =
            "local c = math.abs(tonumber(redis.call('get', KEYS[1]))); local x = math.abs(tonumber(ARGV[1])); if c then if x > c then redis.call('set', KEYS[1], ARGV[1]); return x else return c end else redis.call('set', KEYS[1], ARGV[1]); return x end";

    private final String setMaxAbsScriptSha;

    public MaxAbsScaler(int windowSize, int[] featureIndices, String host, int port, String name) {
        this.windowSize = windowSize;
        this.featureIndices = featureIndices;
        this.host = host;
        this.port = port;
        this.prefix = name + "_max_abs_scaler";

        keys = new String[featureIndices.length];
        channel = new ArrayBlockingQueue<>(1);

        Jedis jedis = new Jedis(host, port);

        Transaction tx = jedis.multi();

        for (int i = 0; i < featureIndices.length; i++) {
            keys[i] = getMaxAbsKey(featureIndices[i]);
            maxAbsValues[i] = Double.MIN_VALUE;
            tx.del(keys[i]);
        }
        tx.exec();
        tx.close();

        this.setMaxAbsScriptSha = jedis.scriptLoad(setMaxAbsScript);

        jedis.close();
    }

    @Override
    public void close() throws Exception {
        super.close();
        channel.offer(new double[0]);
        isRunning = false;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);

        new Thread(
                        () -> {
                            try {
                                syncState();
                            } catch (Exception e) {
                                LOG.info("syncRoutine error: {}", e.getMessage());
                            }
                        })
                .start();
        LOG.info("opened");
    }

    @Override
    public void processElement(
            DenseVector value,
            ProcessFunction<DenseVector, DenseVector>.Context ctx,
            Collector<DenseVector> out)
            throws Exception {
        for (int i = 0; i < featureIndices.length; i++) {
            maxAbsValues[i] = Math.max(Math.abs(value.get(featureIndices[i])), maxAbsValues[i]);
        }
        currSize++;

        // process one record
        if (!bootWindow) {
            for (int i = 0; i < featureIndices.length; i++) {
                value.set(featureIndices[i], value.get(featureIndices[i]) / maxAbsValues[i]);
            }
        }
        out.collect(value);

        if (currSize >= windowSize) {
            double[] minValuesBuffer = this.maxAbsValues.clone();
            currSize = 0;
            channel.put(minValuesBuffer);
        }
    }

    private String getMaxAbsKey(int featureIdx) {
        return prefix + "_abs_" + featureIdx;
    }

    public void syncState() throws Exception {
        Jedis jedis = new Jedis(host, port);
        while (isRunning) {
            double[] maxAbsValuesBuffer = channel.take();
            if (maxAbsValuesBuffer.length == 0) {
                LOG.info("syncRoutine exit");
                break;
            }
            Transaction tx = jedis.multi();
            ;

            for (int i = 0; i < featureIndices.length; i++) {
                tx.evalsha(setMaxAbsScriptSha, 1, keys[i], String.valueOf(maxAbsValuesBuffer[i]));
            }

            List<Object> result = tx.exec();

            boolean flag = true;
            for (int i = 0; i < featureIndices.length; i++) {
                Object minObj = result.get(i);
                if (minObj == null) {
                    flag = false;
                    break;
                }
                maxAbsValuesBuffer[i] = Double.parseDouble((String) minObj);
            }
            tx.close();
            if (!flag) {
                continue;
            }

            this.maxAbsValues = maxAbsValuesBuffer;
            this.bootWindow = false;
        }
        jedis.close();
    }
}
